import pickle
import logging
from collections import defaultdict, Counter
from typing import List, Dict, Any, Set, Optional


class SimilarityMatcher:
    """
    基于倒排索引的高效相似度匹配器（支持类别隔离）
    支持在百万级数据中毫秒级找出最相似记录
    """

    def __init__(self):

        # 存储结构：category -> (inverted_index, candidates, next_id)
        self._data_by_category = defaultdict(lambda: {
            "inverted_index": defaultdict(set),
            "candidates": {},
            "next_id": 0
        })

    def _get_store(self, category: str):
        """获取指定类别的存储结构"""
        return self._data_by_category[category]

    def add_candidate(
            self,
            record: Dict[str, Any],
            record_id: Any = None,
            category: str = "default"
    ) -> Any:
        """
        添加一条候选记录
        :param record: 记录字典
        :param record_id: 可选的自定义 ID
        :param category: 记录所属类别
        :return: 分配的 record_id
        """
        store = self._get_store(category)
        inverted_index = store["inverted_index"]
        candidates = store["candidates"]

        if record_id is None:
            record_id = store["next_id"]
            store["next_id"] += 1

        # 清洗数据：去除排除字段
        filtered_record = {
            k: v for k, v in record.items()
        }

        # 如果记录已存在，先移除旧的倒排索引
        if record_id in candidates:
            old_record = candidates[record_id]
            for k, v in old_record.items():
                key = (k, v)
                if key in inverted_index:
                    inverted_index[key].discard(record_id)

        candidates[record_id] = filtered_record

        # 更新倒排索引
        for k, v in filtered_record.items():
            inverted_index[(k, v)].add(record_id)

        return record_id

    def add_candidates(
            self,
            records: List[Dict[str, Any]],
            category: str = "default"
    ):
        """
        批量添加候选数据
        """
        for record in records:
            self.add_candidate(record, category=category)

    def remove_candidate(self, record_id: Any, category: str = "default"):
        """
        删除一条候选记录
        """
        store = self._get_store(category)
        candidates = store["candidates"]
        inverted_index = store["inverted_index"]

        if record_id not in candidates:
            return

        record = candidates[record_id]
        for k, v in record.items():
            key = (k, v)
            if key in inverted_index:
                inverted_index[key].discard(record_id)

        del candidates[record_id]

    def find_most_similar(
            self,
            current: Dict[str, Any],
            top_k: int = 1,
            min_match_count: int = 1,
            category: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        找出最相似的记录（支持 Top-K 和类别过滤）
        :param current: 当前输入数据
        :param top_k: 返回前 K 个最相似的
        :param min_match_count: 最少匹配属性数
        :param category: 指定查询的类别，None 表示所有类别
        :return: 匹配结果列表
        """
        # 清洗查询数据
        query_items = [
            (k, v) for k, v in current.items()
        ]

        if not query_items:
            return []

        score_counter = Counter()

        if category is not None:
            # 只在指定类别中搜索
            if category not in self._data_by_category:
                return []
            store = self._data_by_category[category]
            inverted_index = store["inverted_index"]
            candidates = store["candidates"]

            for k, v in query_items:
                key = (k, v)
                matching_ids = inverted_index.get(key, set())
                for record_id in matching_ids:
                    if record_id in candidates:  # 确保记录未被删除
                        score_counter[record_id] += 1
        else:
            # 搜索所有类别
            for cat, store in self._data_by_category.items():
                inverted_index = store["inverted_index"]
                candidates = store["candidates"]
                for k, v in query_items:
                    key = (k, v)
                    matching_ids = inverted_index.get(key, set())
                    for record_id in matching_ids:
                        if record_id in candidates:
                            score_counter[record_id] += 1

        # 过滤并排序
        qualified = [
            (rid, count) for rid, count in score_counter.items()
            if count >= min_match_count
        ]
        qualified.sort(key=lambda x: x[1], reverse=True)
        top_k_results = qualified[:top_k]

        results = []
        # 确定从哪个类别取数据（如果 category 指定，则只从该类取）
        categories_to_lookup = [category] if category is not None else self._data_by_category.keys()

        # 构建结果（需找到 record_id 所在的 category）
        for record_id, match_count in top_k_results:
            candidate = None
            found_category = None
            for cat in categories_to_lookup:
                if record_id in self._data_by_category[cat]["candidates"]:
                    candidate = self._data_by_category[cat]["candidates"][record_id]
                    found_category = cat
                    break
            if not candidate:
                continue

            matched_attrs = [
                k for k, v in candidate.items()
                if current.get(k) == v
            ]

            results.append({
                "match_count": match_count,
                "match_rate": round(match_count / len(candidate), 4),
                "record_id": record_id,
                "candidate": candidate,
                "matched_attrs": matched_attrs,
                "found_category": found_category
            })

        return results

    def size(self, category: Optional[str] = None) -> int:
        """
        返回候选集大小
        :param category: 指定类别，None 表示所有类别
        :return: 总记录数
        """
        if category is not None:
            return len(self._data_by_category[category]["candidates"])
        return sum(len(store["candidates"]) for store in self._data_by_category.values())

    def clear(self, category: Optional[str] = None):
        """
        清空数据
        :param category: 指定类别，None 表示清空所有
        """
        if category is not None:
            if category in self._data_by_category:
                del self._data_by_category[category]
        else:
            self._data_by_category.clear()

    def get_categories(self) -> List[str]:
        """获取所有已使用的类别"""
        return list(self._data_by_category.keys())


# 全局 matcher 实例
matcher: Optional[SimilarityMatcher] = SimilarityMatcher()


# ================== Redis 缓存管理 ==================

class RedisCache:
    def __init__(self, host="localhost", port=6379, db=0, key="matcher:v1"):
        import redis
        self.redis_client = redis.Redis(host=host, port=port, db=db, decode_responses=False)
        self.cache_key = key

    def save(self, matcher: SimilarityMatcher):
        """将 matcher 序列化并保存到 Redis"""
        try:
            data = pickle.dumps(matcher)
            self.redis_client.set(self.cache_key, data)
            logging.info(f"✅ 倒排索引已保存到 Redis: {self.cache_key}")
        except Exception as e:
            logging.error(f"❌ 保存失败: {e}")

    def load(self) -> Optional[SimilarityMatcher]:
        """从 Redis 加载 matcher"""
        try:
            data = self.redis_client.get(self.cache_key)
            if data:
                matcher = pickle.loads(data)
                logging.info(f"✅ 从 Redis 加载倒排索引: {self.cache_key}")
                return matcher
            else:
                logging.info("🆕 未找到缓存，将创建新 matcher")
                return None
        except Exception as e:
            logging.error(f"❌ 加载失败: {e}")
            return None
